#!/usr/bin/env python3
"""
多波段预测功能测试脚本
测试训练后的模型是否能进行多波段预测
"""

import torch
import numpy as np
import yaml
import os
from data_utils import load_config
from models import CombinedModel
from datetime import datetime

def test_multi_band_prediction():
    """测试多波段预测功能"""
    print("=" * 50)
    print("测试多波段预测功能")
    print("=" * 50)
    
    # 加载配置
    config = load_config()
    band_list = config['data']['himawari_bands']
    print(f"配置中的波段列表: {band_list}")
    
    # 检查是否有训练好的模型
    model_path = os.path.join(config['training']['save_dir'], 'best_model.pth')
    
    if os.path.exists(model_path):
        print(f"找到训练好的模型: {model_path}")
        
        # 加载模型
        device = torch.device("cpu")  # 使用CPU进行预测
        model = CombinedModel(config).to(device)
        
        try:
            checkpoint = torch.load(model_path, map_location=device)
            model.load_state_dict(checkpoint['model_state_dict'])
            model.eval()
            print(f"模型加载成功，训练轮数: {checkpoint['epoch']}")
        except Exception as e:
            print(f"模型加载失败: {e}")
            print("使用未训练的模型进行测试")
    else:
        print("未找到训练好的模型，使用未训练的模型进行测试")
        device = torch.device("cpu")
        model = CombinedModel(config).to(device)
        model.eval()
    
    # 创建测试数据
    batch_size = 1
    historical_points = config['data']['historical_days'] * 24 * 4  # 15分钟间隔
    num_channels = config['model']['num_channels']
    
    # 模拟场站数据 [batch_size, historical_points, 5]
    station_data = torch.randn(batch_size, historical_points, 5)
    
    # 模拟多波段卫星数据 [batch_size, channels, height, width]
    himawari_data = torch.randn(batch_size, num_channels, 224, 224)
    
    print(f"场站数据形状: {station_data.shape}")
    print(f"卫星数据形状: {himawari_data.shape}")
    
    # 测试预测
    try:
        with torch.no_grad():
            predictions = model(station_data, himawari_data)
            future_points = config['data']['future_hours'] * 4  # 15分钟间隔
            expected_shape = (batch_size, future_points, config['model']['num_output_vars'])
            
            print(f"预测输出形状: {predictions.shape}")
            print(f"期望输出形状: {expected_shape}")
            
            if predictions.shape == expected_shape:
                print("✅ 多波段预测测试通过")
                
                # 显示预测结果统计
                pred_np = predictions.numpy()
                print(f"预测结果统计:")
                print(f"  最小值: {pred_np.min():.4f}")
                print(f"  最大值: {pred_np.max():.4f}")
                print(f"  均值: {pred_np.mean():.4f}")
                print(f"  标准差: {pred_np.std():.4f}")
            else:
                print("❌ 预测输出形状不匹配")
                
    except Exception as e:
        print(f"❌ 预测失败: {e}")
    
    print()

def test_configuration():
    """测试配置是否正确"""
    print("=" * 50)
    print("测试配置")
    print("=" * 50)
    
    config = load_config()
    
    # 检查波段配置
    band_list = config['data']['himawari_bands']
    num_channels = config['model']['num_channels']
    
    print(f"波段数量: {len(band_list)}")
    print(f"模型输入通道数: {num_channels}")
    print(f"历史天数: {config['data']['historical_days']}")
    print(f"预测小时数: {config['data']['future_hours']}")
    print(f"时间分辨率: {config['data']['time_resolution']}")
    
    if len(band_list) == num_channels:
        print("✅ 波段配置正确")
    else:
        print("❌ 波段配置不匹配")
    
    print()

def main():
    """主测试函数"""
    print("气象预测AI系统多波段预测功能测试")
    print("=" * 60)
    
    # 测试配置
    test_configuration()
    
    # 测试多波段预测
    test_multi_band_prediction()
    
    print("=" * 60)
    print("测试完成！")
    print("如果所有测试都通过，说明多波段预测功能正常工作")

if __name__ == "__main__":
    main()
